# set up environment ----
library(data.table)
library(parallel)
library(nnet)
library(RColorBrewer)
library(ggplot2)
library(cowplot)
load("data/analysis_data.RData")
source("R/00-functions.R")
RNGkind("L'Ecuyer-CMRG")
set.seed(1444055996)

# main dataset ----
data_main <- respondent_data[main == 1]
model_main <- multinom(type2 ~ order, data = data_main)
boot <- function(i) {
  multinom(type2 ~ order, data = data_main[sample(.N, .N, replace = TRUE)])
}
boots_main <- mclapply(1:1000, boot, mc.cores = 6)
predicted_probs_main_1 <- do.call(rbind, lapply(boots_main, predict,
  newdata = data.table(order = "Electability Elicited Last"), type = "probs"))
predicted_probs_main_1 <- as.data.table(predicted_probs_main_1)
predicted_probs_main_1[, `:=`(order = "Electability Elicited Last", b = 1:1000)]
predicted_probs_main_0 <- do.call(rbind, lapply(boots_main, predict,
  newdata = data.table(order = "Preferences Elicited Last"), type = "probs"))
predicted_probs_main_0 <- as.data.table(predicted_probs_main_0)
predicted_probs_main_0[, `:=`(order = "Preferences Elicited Last", b = 1:1000)]
type_distributions_main <- rbind(
  predicted_probs_main_1[, .(
    treat = 1,
    Preferences = median(preferences),
    Electability = median(electability),
    `Expected Utility` = median(expected_utility)
  ), b],
  predicted_probs_main_0[, .(
    treat = 0,
    Preferences = median(preferences),
    Electability = median(electability),
    `Expected Utility` = median(expected_utility)
  ), b])

# full dataset ----
data_full <- respondent_data
model_full <- multinom(type2 ~ order, data = data_full)
boot <- function(i) {
  multinom(type2 ~ order, data = data_full[sample(.N, .N, replace = TRUE)])
}
boots_full <- mclapply(1:1000, boot, mc.cores = 6)
predicted_probs_full_1 <- do.call(rbind, lapply(boots_full, predict,
  newdata = data.table(order = "Electability Elicited Last"), type = "probs"))
predicted_probs_full_1 <- as.data.table(predicted_probs_full_1)
predicted_probs_full_1[, `:=`(order = "Electability Elicited Last", b = 1:1000)]
predicted_probs_full_0 <- do.call(rbind, lapply(boots_full, predict,
  newdata = data.table(order = "Preferences Elicited Last"), type = "probs"))
predicted_probs_full_0 <- as.data.table(predicted_probs_full_0)
predicted_probs_full_0[, `:=`(order = "Preferences Elicited Last", b = 1:1000)]
type_distributions_full <- rbind(
  predicted_probs_full_1[, .(
    treat = 1,
    Preferences = median(preferences),
    Electability = median(electability),
    `Expected Utility` = median(expected_utility)
  ), b],
  predicted_probs_full_0[, .(
    treat = 0,
    Preferences = median(preferences),
    Electability = median(electability),
    `Expected Utility` = median(expected_utility)
  ), b])

# ternary plot (main) ----
gridlines_x_intercept = c(.25, .5, .75)
fitted_thetas0 <- type_distributions_main[treat == 0,
  cbind(Preferences, Electability, `Expected Utility`)]
fitted_thetas1 <- type_distributions_main[treat == 1,
  cbind(Preferences, Electability, `Expected Utility`)]
barycentric0 <- data.table(
  x = (fitted_thetas0[, 2] * .5 + fitted_thetas0[, 3]),
  y = (fitted_thetas0[, 2] * .5 * sqrt(3)))
mean_theta_values0 <- apply(fitted_thetas0, 2, mean)
mean_thetas0 <- data.table(
  x = mean_theta_values0[2] * .5 + mean_theta_values0[3],
  y = mean_theta_values0[2] * .5 * sqrt(3))
barycentric1 <- data.table(
  x = (fitted_thetas1[, 2] * .5 + fitted_thetas1[, 3]),
  y = (fitted_thetas1[, 2] * .5 * sqrt(3)))
mean_theta_values1 <- apply(fitted_thetas1, 2, mean)
mean_thetas1 <- data.table(
  x = mean_theta_values1[2] * .5 + mean_theta_values1[3],
  y = mean_theta_values1[2] * .5 * sqrt(3))
r <- .05
h1 <- .09
h2 <- h1 * .6
h3 <- .02
w1 <- .25
gsize <- .5
g0 <- ggplot(barycentric0, aes(x, y)) +
  geom_density2d_filled(
    data = barycentric0,
    aes(fill = ..level.., color = ..level..),
    contour_var = "ndensity",
    breaks = exp(seq(log(.025), log(.975), length.out = 70)) / exp(log(.975)),
    n = 200) +
  geom_density2d_filled(
    data = barycentric1,
    aes(fill = ..level.., color = ..level..),
    contour_var = "ndensity",
    breaks = exp(seq(log(.025), log(.975), length.out = 70)) / exp(log(.975)),
    n = 200) +
  annotate(geom = "segment", x = 0, xend = 0.5, y = 0, yend = 0.5 * sqrt(3),
    color = "black", size = gsize) +
  annotate(geom = "segment", x = 1, xend = 0.5, y = 0, yend = 0.5 * sqrt(3),
    color = "black", size = gsize) +
  annotate(geom = "segment", x = 0, xend = 1, y = 0, yend = 0,
    color = "black", size = gsize) +
  annotate(geom = "segment", x = .25, xend = .625, y = 0, yend = 0.6495191,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .5, xend = 0.750, y = 0, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .75, xend = 0.875, y = 0, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .25, xend = 0.125, y = 0, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .5, xend = 0.250, y = 0, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .75, xend = 0.375, y = 0, yend = 0.6495191,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.125, xend = 0.875, y = 0.2165064, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.250, xend = 0.750, y = 0.4330127, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.375, xend = 0.625, y = 0.6495191, yend = 0.6495191,
    color = "black", size = gsize/2) +
  geom_point(data = mean_thetas0, color = "white", size = 2) +
  geom_point(data = mean_thetas1, color = "black", size = 2) +
  scale_fill_manual(values = paste0("gray", 100:31)) +#brewer.pal(9, "Greys")) +
  scale_color_manual(values = paste0("gray", 100:31)) +#brewer.pal(9, "Greys")) +
  theme_bw() +
  theme(
    axis.title = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    axis.line = element_blank(),
    panel.grid = element_blank(),
    panel.border = element_blank(),
    legend.position = "none") +
  annotate(
    geom = "text",
    x = 1/2 + h1 * (0),
    y = 0 + h1 * (-1),
    angle = 0,
    label = "Preferences") +
  annotate(
    geom = "text",
    x = 1/4 + h1 * (-4/5),
    y = sqrt(3)/4 + h1 * (3/5),
    angle = 60,
    label = "Electability") +
  annotate(
    geom = "text",
    x = 3/4 + h1 * (4/5),
    y = sqrt(3)/4 + h1 * (3/5),
    angle = -60,
    label = "Expected Utility") +
  annotate(geom = "segment",
    # origin + normal vector to axis (for text) + parallel vector to axis (for arrow)
    x = 1/4 + h2 * (-4/5) - w1 * (1/2),
    y = sqrt(3)/4 + h2 * (3/5) - w1 * (sqrt(3)/2),
    xend = 1/4 + h2 * (-4/5) + w1 * (1/2),
    yend = sqrt(3)/4 + h2 * (3/5) + w1 * (sqrt(3)/2),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches"))) +
  annotate(geom = "segment",
    x = 1/2 + h2 * (0) - w1 * (-1),
    y = 0 + h2 * (-1) - w1 * (0),
    xend = 1/2 + h2 * (0) + w1 * (-1),
    yend = 0 + h2 * (-1) + w1 * (0),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches")))  +
  annotate(geom = "segment",
    x = 3/4 + h2 * (4/5) - w1 * (1/2),
    y = sqrt(3)/4 + h2 * (3/5) - w1 * (-sqrt(3)/2),
    xend = 3/4 + h2 * (4/5) + w1 * (1/2),
    yend = sqrt(3)/4 + h2 * (3/5) + w1 * (-sqrt(3)/2),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches"))) +
  annotate(
    geom = "text",
    x = .75 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "25%",
    size = 3) +
  annotate(
    geom = "text",
    x = .50 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = .25 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .75 + h3 * (-4/5),
    y = sqrt(3)/2 * .75 + h3 * (3/5),
    angle = 60,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .5 + h3 * (-4/5),
    y = sqrt(3)/2 * .5 + h3 * (3/5),
    angle = 60,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .25 + h3 * (-4/5),
    y = sqrt(3)/2 * .25 + h3 * (3/5),
    angle = 60,
    label = "25%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .75 + h3 * (4/5),
    y = sqrt(3)/2 * .25 + h3 * (3/5),
    angle = -60,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .5 + h3 * (4/5),
    y = sqrt(3)/2 * .5 + h3 * (3/5),
    angle = -60,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .25 + h3 * (4/5),
    y = sqrt(3)/2 * .75 + h3 * (3/5),
    angle = -60,
    label = "25%",
    size = 3) +
  coord_equal(clip = "off", xlim = c(-.0, 1.0), ylim = c(-.0, sqrt(3)/2+.0),
    expand = FALSE) +
  theme(plot.margin = unit(c(1,1,1,0), "lines")) +
  annotate(
    geom = "text",
    x = .5,
    y = .37 * sqrt(3)/2,
    label = "Electability Last",
    size = 3.2) +
  annotate(
    geom = "text",
    x = .4,
    y = .08 * sqrt(3)/2,
    label = "Preferences Last",
    size = 3.2)

cairo_pdf("figures/ternary-main-type2.pdf", width = 3.7, height = 3.7)
ggdraw() +
  draw_plot(g0, x = .05 / 2, y = 0, width = .95, height = .95) +
  draw_label("Elicitation Order Affects Type Distribution", 0.5, 0.93, size = 12)
dev.off()

# ternary plot (full) ----
gridlines_x_intercept = c(.25, .5, .75)
fitted_thetas0 <- type_distributions_full[treat == 0,
  cbind(Preferences, Electability, `Expected Utility`)]
fitted_thetas1 <- type_distributions_full[treat == 1,
  cbind(Preferences, Electability, `Expected Utility`)]
barycentric0 <- data.table(
  x = (fitted_thetas0[, 2] * .5 + fitted_thetas0[, 3]),
  y = (fitted_thetas0[, 2] * .5 * sqrt(3)))
mean_theta_values0 <- apply(fitted_thetas0, 2, mean)
mean_thetas0 <- data.table(
  x = mean_theta_values0[2] * .5 + mean_theta_values0[3],
  y = mean_theta_values0[2] * .5 * sqrt(3))
barycentric1 <- data.table(
  x = (fitted_thetas1[, 2] * .5 + fitted_thetas1[, 3]),
  y = (fitted_thetas1[, 2] * .5 * sqrt(3)))
mean_theta_values1 <- apply(fitted_thetas1, 2, mean)
mean_thetas1 <- data.table(
  x = mean_theta_values1[2] * .5 + mean_theta_values1[3],
  y = mean_theta_values1[2] * .5 * sqrt(3))
r <- .05
h1 <- .09
h2 <- h1 * .6
h3 <- .02
w1 <- .25
gsize <- .5
g0 <- ggplot(barycentric0, aes(x, y)) +
  geom_density2d_filled(
    data = barycentric0,
    aes(fill = ..level.., color = ..level..),
    contour_var = "ndensity",
    breaks = exp(seq(log(.025), log(.975), length.out = 70)) / exp(log(.975)),
    n = 200) +
  geom_density2d_filled(
    data = barycentric1,
    aes(fill = ..level.., color = ..level..),
    contour_var = "ndensity",
    breaks = exp(seq(log(.025), log(.975), length.out = 70)) / exp(log(.975)),
    n = 200) +
  annotate(geom = "segment", x = 0, xend = 0.5, y = 0, yend = 0.5 * sqrt(3),
    color = "black", size = gsize) +
  annotate(geom = "segment", x = 1, xend = 0.5, y = 0, yend = 0.5 * sqrt(3),
    color = "black", size = gsize) +
  annotate(geom = "segment", x = 0, xend = 1, y = 0, yend = 0,
    color = "black", size = gsize) +
  annotate(geom = "segment", x = .25, xend = .625, y = 0, yend = 0.6495191,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .5, xend = 0.750, y = 0, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .75, xend = 0.875, y = 0, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .25, xend = 0.125, y = 0, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .5, xend = 0.250, y = 0, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = .75, xend = 0.375, y = 0, yend = 0.6495191,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.125, xend = 0.875, y = 0.2165064, yend = 0.2165064,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.250, xend = 0.750, y = 0.4330127, yend = 0.4330127,
    color = "black", size = gsize/2) +
  annotate(geom = "segment", x = 0.375, xend = 0.625, y = 0.6495191, yend = 0.6495191,
    color = "black", size = gsize/2) +
  geom_point(data = mean_thetas0, color = "white", size = 2) +
  geom_point(data = mean_thetas1, color = "black", size = 2) +
  scale_fill_manual(values = paste0("gray", 100:31)) +#brewer.pal(9, "Greys")) +
  scale_color_manual(values = paste0("gray", 100:31)) +#brewer.pal(9, "Greys")) +
  theme_bw() +
  theme(
    axis.title = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    axis.line = element_blank(),
    panel.grid = element_blank(),
    panel.border = element_blank(),
    legend.position = "none") +
  annotate(
    geom = "text",
    x = 1/2 + h1 * (0),
    y = 0 + h1 * (-1),
    angle = 0,
    label = "Preferences") +
  annotate(
    geom = "text",
    x = 1/4 + h1 * (-4/5),
    y = sqrt(3)/4 + h1 * (3/5),
    angle = 60,
    label = "Electability") +
  annotate(
    geom = "text",
    x = 3/4 + h1 * (4/5),
    y = sqrt(3)/4 + h1 * (3/5),
    angle = -60,
    label = "Expected Utility") +
  annotate(geom = "segment",
    # origin + normal vector to axis (for text) + parallel vector to axis (for arrow)
    x = 1/4 + h2 * (-4/5) - w1 * (1/2),
    y = sqrt(3)/4 + h2 * (3/5) - w1 * (sqrt(3)/2),
    xend = 1/4 + h2 * (-4/5) + w1 * (1/2),
    yend = sqrt(3)/4 + h2 * (3/5) + w1 * (sqrt(3)/2),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches"))) +
  annotate(geom = "segment",
    x = 1/2 + h2 * (0) - w1 * (-1),
    y = 0 + h2 * (-1) - w1 * (0),
    xend = 1/2 + h2 * (0) + w1 * (-1),
    yend = 0 + h2 * (-1) + w1 * (0),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches")))  +
  annotate(geom = "segment",
    x = 3/4 + h2 * (4/5) - w1 * (1/2),
    y = sqrt(3)/4 + h2 * (3/5) - w1 * (-sqrt(3)/2),
    xend = 3/4 + h2 * (4/5) + w1 * (1/2),
    yend = sqrt(3)/4 + h2 * (3/5) + w1 * (-sqrt(3)/2),
    size = .25,
    arrow = arrow(type = "closed", angle = 15, length = unit(.125, "inches"))) +
  annotate(
    geom = "text",
    x = .75 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "25%",
    size = 3) +
  annotate(
    geom = "text",
    x = .50 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = .25 + h3 * (0),
    y = 0 + h3 * (-1),
    angle = 0,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .75 + h3 * (-4/5),
    y = sqrt(3)/2 * .75 + h3 * (3/5),
    angle = 60,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .5 + h3 * (-4/5),
    y = sqrt(3)/2 * .5 + h3 * (3/5),
    angle = 60,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 * .25 + h3 * (-4/5),
    y = sqrt(3)/2 * .25 + h3 * (3/5),
    angle = 60,
    label = "25%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .75 + h3 * (4/5),
    y = sqrt(3)/2 * .25 + h3 * (3/5),
    angle = -60,
    label = "75%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .5 + h3 * (4/5),
    y = sqrt(3)/2 * .5 + h3 * (3/5),
    angle = -60,
    label = "50%",
    size = 3) +
  annotate(
    geom = "text",
    x = 1/2 + 1/2 * .25 + h3 * (4/5),
    y = sqrt(3)/2 * .75 + h3 * (3/5),
    angle = -60,
    label = "25%",
    size = 3) +
  coord_equal(clip = "off", xlim = c(-.0, 1.0), ylim = c(-.0, sqrt(3)/2+.0),
    expand = FALSE) +
  theme(plot.margin = unit(c(1,1,1,0), "lines")) +
  annotate(
    geom = "text",
    x = .5,
    y = .37 * sqrt(3)/2,
    label = "Electability Last",
    size = 3.2) +
  annotate(
    geom = "text",
    x = .4,
    y = .08 * sqrt(3)/2,
    label = "Preferences Last",
    size = 3.2)

cairo_pdf("figures/ternary-full-type2.pdf", width = 3.7, height = 3.7)
ggdraw() +
  draw_plot(g0, x = .05 / 2, y = 0, width = .95, height = .95) +
  draw_label("Elicitation Order Affects Type Distribution\n(All Respondents)", 0.5, 0.93, size = 12)
dev.off()

# save ----
save(boots_main, file = "results/boots_main-type2.RData")
save(predicted_probs_main_1, file = "results/predicted_probs_main_1-type2.RData")
save(predicted_probs_main_0, file = "results/predicted_probs_main_0-type2.RData")
save(boots_full, file = "results/boots_full-type2.RData")
